package vulnerability

import (
	"fmt"
	fingerprint "github.com/yqcs/fingerscan"
	"net"
	"prismx_cli/core/models"
	"prismx_cli/core/plugins"
	"prismx_cli/utils/arr"
	"prismx_cli/utils/logger"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Query struct {
	App      string
	Title    string
	Header   string
	Body     string
	Uri      string
	Port     int
	Protocol string
	Symbol   string
}

func Exploit(name, scheme, ip string, port int, payload string, isRed bool, timeout time.Duration) (res models.VulResult) {
	for _, item := range plugins.GetAllAppVulList() {
		if item.AppVul.Meta.Name == name {
			if item.PluginType == "Service" {
				if isRed && payload != "" {
					if scanRes := AddReflect(item.AppVul.Meta.Steps.ExploitSteps.ExploitGo, scheme, ip, port, payload, timeout); scanRes != nil {
						res = scanRes[0].Interface().(models.VulResult)
					}
				} else {
					if scanRes := AddReflect(item.AppVul.Meta.Steps.VerifySteps.VerifyGo, scheme, ip, port, timeout); scanRes != nil {
						res = scanRes[0].Interface().(models.VulResult)
					}
				}
				return
			} else {
				//如果服务不是http就给替换为http，仅生效于web插件，examples: 服务是docker，替换成http
				if scheme != "http" && scheme != "https" {
					scheme = "http"
				}
				res = HttpVerify(scheme+"://"+net.JoinHostPort(ip, strconv.Itoa(port)), item.AppVul.Meta, payload, isRed, timeout)
				return
			}
		}
	}
	return res
}

// Verify 匹配指纹模式
// 参数 指纹、taskId 超时时间
func Verify(r *fingerprint.AppFinger, timeout time.Duration) {
	wg := sync.WaitGroup{}
	//如果匹配到插件库就直接选择并利用，没有匹配到指纹信息说明不配使用ai模块
	for _, vulItem := range plugins.GetAllAppVulList() {
		wg.Add(1)
		go func(item plugins.Vul) {

			defer wg.Done()                            //结束时关闭
			query := ParseQueryJson(item.AppVul.Query) //解析查询语法
			isFlag := false
			//进行规则匹配
			switch {
			//如果App等于* 或者网站的技术包含该item的技术，并且query等于该技术，则通过
			case query.App == "*" || (query.App != "" && arr.FuzzContainToLower(r.WebApp.App, query.App)):
				isFlag = true
			case query.Title != "" && strings.Contains(strings.ToLower(r.WebApp.Title), strings.ToLower(query.Title)):
				isFlag = true
			case query.Header != "" && strings.Contains(strings.ToLower(r.WebApp.Header), strings.ToLower(query.Header)):
				isFlag = true
			case query.Body != "" && strings.Contains(strings.ToLower(r.WebApp.Body), strings.ToLower(query.Body)):
				isFlag = true
			case query.Uri != "" && strings.Contains(r.Uri, query.Uri):
				isFlag = true
			case query.Port != 0 && r.Port == query.Port:
				isFlag = true
			case query.Protocol != "" && strings.Contains(r.Service, strings.ToLower(query.Protocol)):
				isFlag = true
			}

			if isFlag {

				var res models.VulResult

				if item.PluginType == "Service" {
					scanRes := AddReflect(item.AppVul.Meta.Steps.VerifySteps.VerifyGo, r.Service, r.IP, r.Port, timeout)
					if scanRes != nil {
						res = scanRes[0].Interface().(models.VulResult)
					}
				} else {
					//如果服务不是http就给替换为http，仅生效于web插件，examples: 服务是docker，替换成http
					if r.Service != "http" && r.Service != "https" {
						r.Service = "http"
					}
					res = HttpVerify(r.Service+"://"+r.Uri, item.AppVul.Meta, "", false, timeout)
				}

				if res.State {
					leveTxt := "Nil"
					if item.AppVul.Meta.Level == 1 {
						leveTxt = "Info"
					} else if item.AppVul.Meta.Level == 2 {
						leveTxt = "Low"
					} else if item.AppVul.Meta.Level == 3 {
						leveTxt = "Middle"
					} else if item.AppVul.Meta.Level == 4 {
						leveTxt = "High"
					} else if item.AppVul.Meta.Level == 5 {
						leveTxt = "Serious"
					}
					msg := fmt.Sprintf("Vuln:%s\n\rTarget: %s\n\rLeve: %s\n\rExp: %t\n\rType: %s\n\rPayload:\n\r————\n\r%s\n\r————", item.AppVul.Meta.Name, strings.Join(item.AppVul.Meta.Tags, ""), leveTxt, item.AppVul.Meta.Available, strings.Join(item.AppVul.Meta.Tags, ""), res.Request)
					logger.ScanSuccess(msg)
				}
			}
		}(vulItem)
	}
	wg.Wait()
}

// ParseQueryJson 解析定位语法
// app、title、header、body、uri、protocol、port
func ParseQueryJson(s string) (query Query) {
	//表达式包含App
	if strings.Contains(s, "app:") {
		query.App = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "title:") {
		query.Title = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "header:") {
		query.Header = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "body:") {
		query.Body = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "uri:") {
		query.Uri = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "protocol:") {
		query.Protocol = s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
	}
	if strings.Contains(s, "port:") {
		p := s[strings.Index(s, `"`)+1 : strings.LastIndex(s, `"`)]
		port, err := strconv.Atoi(p)
		if err != nil {
			query.Port = 0
		}
		query.Port = port
	}
	return query
}

// AddReflect 反射扫描函数
func AddReflect(app any, infos ...any) []reflect.Value {
	if app == nil {
		return nil
	}
	f := reflect.ValueOf(app)
	//判断方法参数与传入是否匹配
	if len(infos) != f.Type().NumIn() {
		logger.Error("The number of infos is not adapted")
		return nil
	}
	in := make([]reflect.Value, len(infos))
	for k, info := range infos {
		in[k] = reflect.ValueOf(info)
	}
	return f.Call(in)
}
